package com.aispeech.asr.comm.data.pojo;

import com.aispeech.asr.comm.constants.Progress;
import com.aispeech.asr.comm.utils.BeanUtils;
import com.aispeech.jarpc.op.post.AuxInfo;
import com.aispeech.jarpc.op.post.Onebest;
import com.aispeech.jarpc.op.post.Segment;
import com.aispeech.jarpc.op.post.v2.Onebest_v2;
import com.aispeech.jarpc.op.post.v2.Segment_v2;
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.databind.PropertyNamingStrategy;
import com.fasterxml.jackson.databind.annotation.JsonNaming;
import lombok.ToString;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.bson.types.ObjectId;
import org.springframework.data.mongodb.core.mapping.Document;
import org.springframework.data.mongodb.core.mapping.MongoId;

import java.time.Instant;
import java.util.*;

@Slf4j
@ToString
@Document("lasr_offline_tasks")
@JsonNaming(PropertyNamingStrategy.SnakeCaseStrategy.class)
public final class TaskDoc {
    int progress;
    List<Map<String, Object>> result;
    @JsonInclude(JsonInclude.Include.NON_NULL)
    Map[] asr;
    @JsonInclude(JsonInclude.Include.NON_NULL)
    List<Map<String, Object>> segment;
    @MongoId
    private String id;
    private String productId;
    @JsonInclude(JsonInclude.Include.NON_NULL)
    private int priority = 0;
    private String hostname = null;
    private long update_t;
    @JsonIgnoreProperties
    private Instant finished_t;
    private Map<String, Object> request;
    @JsonInclude(JsonInclude.Include.NON_NULL)
    private Map<String, Long> metrics;
    @JsonInclude(JsonInclude.Include.NON_NULL)
    private Map<String, Object> error;
    @JsonInclude(JsonInclude.Include.NON_NULL)
    private String parentId;
    @JsonInclude(JsonInclude.Include.NON_NULL)
    private List<TaskDoc> channels;
    @JsonInclude(JsonInclude.Include.NON_DEFAULT)
    private int channelIndex = 0;
    private String res;
    private String resVersion;

    public TaskDoc() {
        metrics = new HashMap<String, Long>() {{
            put("received_t", System.currentTimeMillis());
        }};
        result = new ArrayList<>();
    }

    public static TaskDoc of(String productId, Map<String, Object> request) {
        TaskDoc task = new TaskDoc();
        task.setId(new ObjectId().toString());
        task.setProductId(productId);
        task.setProgress(Progress.IDLE);
        task.setRequest(request);

        return task;
    }

    public static TaskDoc mergeAnnotation(TaskDoc taskDoc, AnnoTaskDoc annoTaskDoc) {
        if (annoTaskDoc != null) {
            List<LabelSentence> labelSentences = annoTaskDoc.getLabeled();
            int len = labelSentences.size();

            if (len <= taskDoc.getResult().size()) {
                List<Map<String, Object>> results = taskDoc.getResult();
                for (int i = 0; i < len; i++) {
                    Map<String, Object> row = results.get(i);
                    LabelSentence labelSentence = labelSentences.get(i);

                    row.put("bg", labelSentence.getBg());
                    row.put("ed", labelSentence.getEd());
                    row.put("speaker", labelSentence.getSpeaker());
                    row.put("onebest", labelSentence.getOnebest());
                }
            }
        }

        if ((int) taskDoc.getRequest().getOrDefault("debug", 0) != 2) {
            taskDoc.setAsr(null);
        }

        if ((int) taskDoc.getRequest().getOrDefault("useMetrics", 0) != 1) {
            taskDoc.setMetrics(null);
        }

        taskDoc.setParentId(null);
        taskDoc.setChannelIndex(0);
        taskDoc.setError(null);

        Map<String, Object> req = new HashMap<>();
        taskDoc.getRequest().forEach((key, value) -> {
            if (!("debug".equals(key) || "useMetrics".equals(key))) {
                req.put(BeanUtils.humpToSnake(key), value);
            }
        });

        taskDoc.setRequest(req);

        return taskDoc;
    }

    public static TaskDoc multiChannelResult(TaskDoc taskDoc, List<TaskDoc> subTaskDocs) {
        subTaskDocs.sort(Comparator.comparing(TaskDoc::getChannelIndex));

        List<TaskDoc> channels = new ArrayList<>();
        for (int i = 0; i < subTaskDocs.size(); i++) {
            channels.add(mergeAnnotation(subTaskDocs.get(i), null));
        }

        taskDoc.setError(null);
        taskDoc.setChannels(channels);

        return taskDoc;
    }

    public List<Map<String, Object>> getSegment() {
        return segment;
    }

    public void setSegment(List<Segment_v2> segments) {
        this.segment = new ArrayList<>();

        segments.forEach(segment -> {
            Map<String, Object> row = new HashMap<>();

            row.put("bg", segment.getBg());
            row.put("ed", segment.getEd());
            row.put("onebest", segment.getOnebest());
            row.put("speaker", segment.getSpeaker());

            if (StringUtils.isNotEmpty(segment.getWp())) row.put("wp", segment.getWp());

            this.segment.add(row);
        });
    }

    public Instant getFinished_t() {
        return finished_t;
    }

    public void setFinished_t(Instant finished_t) {
        this.finished_t = finished_t;
    }

    public Map[] getAsr() {
        return asr;
    }

    public void setAsr(Map[] asr) {
        this.asr = asr;
    }

    public String getRes() {
        return res;
    }

    public void setRes(String res) {
        this.res = res;
    }

    public String getResVersion() {
        return resVersion;
    }

    public void setResVersion(String resVersion) {
        this.resVersion = resVersion;
    }

    public String getParentId() {
        return parentId;
    }

    public void setParentId(String parentId) {
        this.parentId = parentId;
    }

    public int getChannelIndex() {
        return channelIndex;
    }

    public void setChannelIndex(int channelIndex) {
        this.channelIndex = channelIndex;
    }

    public String getProductId() {
        return productId;
    }

    public void setProductId(String productId) {
        this.productId = productId;
    }

    public int getProgress() {
        return progress;
    }

    public void setProgress(int progress) {
        this.progress = progress;
    }

    public List<Map<String, Object>> getResult() {
        return result;
    }

    public void setResult(List<Onebest_v2> result) {
        this.result = new ArrayList<>();

        result.forEach(onebest -> {
            Map<String, Object> row = new HashMap<>();

            row.put("bg", onebest.getBg());
            row.put("ed", onebest.getEd());
            row.put("onebest", onebest.getOnebest());
            row.put("speaker", onebest.getSpeaker());

            AuxInfo auxInfo = onebest.getAuxinfo();
            if (auxInfo != null) {
                Map<String, Object> auxMap = new HashMap<>();

                auxMap.put("emotion", auxInfo.getEmotion());
                auxMap.put("age", auxInfo.getAge());
                auxMap.put("gender", auxInfo.getGender());

                row.put("auxinfo", auxMap);
            }

            this.result.add(row);
        });
    }

    public String getId() {
        return id;
    }

    public void setId(String id) {
        this.id = id;
    }

    public int getPriority() {
        return priority;
    }

    public void setPriority(int priority) {
        this.priority = priority;
    }

    public String getHostname() {
        return hostname;
    }

    public void setHostname(String hostname) {
        this.hostname = hostname;
    }

    public long getUpdate_t() {
        return update_t;
    }

    public void setUpdate_t(long update_t) {
        this.update_t = update_t;
    }

    public Map<String, Object> getRequest() {
        return request;
    }

    public void setRequest(Map<String, Object> request) {
        this.request = request;
        String appId = (String) request.get("appId");

        if (StringUtils.isNotEmpty(appId))
            productId = appId;
    }

    public Map<String, Long> getMetrics() {
        return metrics;
    }

    public void setMetrics(Map<String, Long> metrics) {
        this.metrics = metrics;
    }

    public void addMetric(String type) {
        if (metrics == null) {
            metrics = new HashMap<>();
        }

        metrics.put(type, System.currentTimeMillis());
    }

    public void addMetric(String type, long ts) {
        if (metrics == null) {
            metrics = new HashMap<>();
        }

        metrics.put(type, ts);
    }

    public Map<String, Object> getError() {
        return error;
    }

    public void setError(Map<String, Object> error) {
        this.error = error;
    }

    public List<TaskDoc> getChannels() {
        return channels;
    }

    public void setChannels(List<TaskDoc> channels) {
        this.channels = channels;
    }
}
